import requests
import re
import sys
from Cryptodome.Cipher import AES
sys.modules['Crypto'] = sys.modules['Cryptodome']

# ====== 配置区域 ======
m3u8_url = "https://www.sikiedu.com/hls/60087/stream/shd/INWi0hPTkMDrOZxhfF1dV6RQH7KTm3P2.m3u8?protocol=https"
output_file = "final_video.mp4"

# 在这里填入你从浏览器 Network 面板复制的请求头（保持字典格式）
headers = {
    "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) "
                  "AppleWebKit/537.36 (KHTML, like Gecko) "
                  "Chrome/139.0.0.0 Safari/537.36",
    "Cookie": "_ga=GA1.1.592550214.1749124105; 53revisit=1749124104880; 53gid2=13513658086000; REMEMBERME=Qml6XFVzZXJcQ3VycmVudFVzZXI6ZFhObGNsODRiRzR3Tm10dFlYaEFaV1IxYzI5b2J5NXVaWFE9OjE3ODU0NjE5MDY6MzZlMDBkOTNhYzM3OTA2NzgyNTEyYWRjN2RkNzIzMWY3ZjUzYTA2YWVkNjY1ZjRmNDZmNDZiMGEwZjJkMjI5Zg%3D%3D; Hm_lvt_7b741b98bcf0532f9add944e2593c8af=1753925802,1754354928; _clck=1s42oov%7C2%7Cfy7%7C0%7C1982; _uetvid=fd8eff30420211f0a982ef8d7a8b55ef; _ga_C0BX4HW2XZ=GS2.1.s1754354927$o3$g1$t1754356736$j60$l0$h0; 53kf_72696701_from_host=www.sikiedu.com; uuid_53kf_72696701=9aef39f1a388ca99e16c771d34f98a59; 53kf_72696701_land_page=https%253A%252F%252Fwww.sikiedu.com%252Fcourse%252F926; kf_72696701_land_page_ok=1; online-uuid=3E8DBC27-B89B-18FD-4EBB-B1D2C2870762; 53kf_72696701_keyword=https%3A%2F%2Fwww.sikiedu.com%2Fpage%2Funitypath; sourceUrl=%E7%9B%B4%E6%8E%A5%E8%AE%BF%E9%97%AE; sl-session=BxJoPvdmsmghwj7EkOVQ4g==; PHPSESSID=1pikhiao6ba0befpaou2v32609; landUrl=https://www.sikiedu.com/page/unitypath",
    # 其他的header（如 Referer、Authorization 等）也可以加
}
# =====================

def download_m3u8(m3u8_url, output_file):
    # 请求 m3u8 文件
    m3u8_content = requests.get(m3u8_url, headers=headers).text

    # 匹配 key 信息
    key = None
    iv = None
    key_url = None

    key_match = re.search(r'#EXT-X-KEY:METHOD=AES-128,URI="(.*?)"(?:,IV=(.*))?', m3u8_content)
    if key_match:
        key_url = key_match.group(1)
        iv = key_match.group(2)
        if iv:
            iv = bytes.fromhex(iv.replace("0x", ""))
        print(f"[+] 发现加密视频，Key 地址: {key_url}")

        # 下载 key
        key = requests.get(key_url, headers=headers).content
        print(f"[+] Key 长度: {len(key)} 字节")

    # 提取 ts 分片
    ts_urls = [line.strip() for line in m3u8_content.split("\n") if line and not line.startswith("#")]
    print(f"[+] 共找到 {len(ts_urls)} 个分片")

    # 拼接成绝对路径
    base_url = m3u8_url.rsplit("/", 1)[0] + "/"
    ts_urls = [url if url.startswith("http") else base_url + url for url in ts_urls]

    # 下载并解密
    with open(output_file, "wb") as f_out:
        for i, ts_url in enumerate(ts_urls):
            ts_data = requests.get(ts_url, headers=headers).content
            if key:
                cipher = AES.new(key, AES.MODE_CBC, iv if iv else ts_data[:16])
                ts_data = cipher.decrypt(ts_data)
            f_out.write(ts_data)
            print(f"[{i+1}/{len(ts_urls)}] 已处理 {ts_url}")

    print(f"[+] 视频保存完成: {output_file}")

if __name__ == "__main__":
    download_m3u8(m3u8_url, output_file)
